#! /usr/bin/env python3

import os
import subprocess as subp
from itertools import product
import re

gnnadv_filt = ["amazon0601", "OVCAR-8H", "amazon0505", "SW-620H", "Yeast", "com-amazon"]
Filter = [
   "FacebookPagePage"
]
Filter += ["GNNAdvData."+x for x in os.listdir("data/GNNAdvData/processed") if x not in gnnadv_filt]
Filter += ["Twitch."+x for x in os.listdir("data/Twitch")]
Filter += ["Planetoid."+x for x in os.listdir("data/Planetoid")]
Filter += ["Coauthor."+x for x in os.listdir("data/Coauthor")]
Filter += ["PPI."+str(x) for x in range(0,20)]

BigFuck = [ "Reddit" ]
BigFuck += ["SNAPDataset."+x for x in os.listdir("data/SNAPDataset")]
BigFuck += ["GNNAdvData."+x for x in gnnadv_filt]

def gen_file_path(mtx_folder, fname):
   parts = fname.split('.')
   if parts[0] in ["Mtx", "GNNAdvData"]:
      return os.path.join(mtx_folder, parts[0], parts[-1]+'.mtx')
   elif parts[0] in ["PPI"]:
      return os.path.join(mtx_folder, parts[0], parts[0]+"_"+parts[-1]+".mtx")
   else:
      return os.path.join(mtx_folder, *parts, parts[-1]+'.mtx')

def gen_sample_path(mtx_folder, fname):
   return os.path.join(mtx_folder, fname)

def run_all_mtx(mtx_folder, dense_range, iterations = 100):
   files = os.listdir(mtx_folder) if len(Filter) == 0 else [x+'.mtx' for x in Filter]
   for f in files:
      for dim in dense_range:
         out = subp.check_output([
            "./build/torch_grad_test/spmm_test",
            os.path.join(mtx_folder, f),
            str(dim), str(iterations), '1'
         ]).strip().split("\n")
         error = out[-1]
         times = [x for x in out if not re.match(r'my time|cusparse time', x) is None]
         my_time = float(times[0].split(' ')[3])
         base_time = float(times[1].split(' ')[3])
         print("{} dim[{}] : time {} speedup {} {}".format(f,dim,my_time,base_time/my_time,error))


class MyConfig(object):
   @classmethod
   def _find_baseline_time(cls, s:str):
      if not re.match(r'cusparse throughput', s) is None:
         return float(s.split(' ')[2])
      else:
         return -1.
   
   @classmethod
   def _find_my_time(cls, s:str):
      if not re.match(r'heads', s) is None:
         return float(s.split(' ')[2][3:-1])
      else:
         return -1.

   @classmethod
   def _get_config_string(cls, s:str):
      return s.split(' ')[0]

   @classmethod
   def _launch_args(cls, file, dim, iterations):
      return [
            "/home/limingyi/gnn-workspace/build/torch_grad_test/spmm_test",
            file,
            '1', str(dim), str(iterations), '2', '1'
         ]


class ManscriptConfig(object):
   @classmethod
   def _find_baseline_time(cls, s:str):
      if not re.match(r'cusparse time', s) is None:
         return float(s.split(' ')[3])
      else:
         return -1.
   
   @classmethod
   def _find_my_time(cls, s:str):
      if not re.match(r'spwalk', s) is None:
         return float(s.split(' ')[2][3:])
      else:
         return -1.

   @classmethod
   def _get_config_string(cls, s:str):
      return s.split(' ')[0]

   @classmethod
   def _launch_args(cls, file, dim, iterations):
      return [
            "/home/limingyi/gnn-workspace/build/torch_grad_test/spmm_test",
            file,
            '1', str(dim), str(iterations), '2', '1'
         ]


class GeSpMMConfig(object):
   @classmethod
   def _find_baseline_time(cls, s:str):
      if not re.match(r"\[Cusparse\]", s) is None:
         return float(s.split(' ')[2])
      else:
         return -1.
   
   @classmethod
   def _find_my_time(cls, s:str):
      if not re.match(r"\[GE\-SpMM\]", s) is None:
         return float(s.split(' ')[2])
      else:
         return -1.
 
   @classmethod
   def _get_config_string(cls, s:str):
      return s.split(' ')[0]

   @classmethod
   def _launch_args(cls, file, dim, iterations):
      return ['/home/limingyi/dgSPARSE-Library/example/ge-spmm/spmm.out',
              file, str(dim), str(iterations)]


class GeSpMMMotiveConfig(object):
   baseline_key = "cusparse"
   conf_keys = [f"[Alg:{x}]" for x in range(10)]

   @classmethod
   def _find_baseline_time(cls, s:str):
      if not re.match(r"\[Cusparse\]", s) is None:
         return float(s.split(' ')[5])
      else:
         return -1.
   
   @classmethod
   def _find_my_time(cls, s:str):
      if not re.match(r"\[GE\-SpMM\]", s) is None:
         return float(s.split(' ')[5])
      else:
         return -1.
 
   @classmethod
   def _get_config_string(cls, s:str):
      return s.split(' ')[0][9:]

   @classmethod
   def _launch_args(cls, file, dim, iterations):
      return ['/home/limingyi/dgSPARSE-Library/example/ge-spmm/spmm.out',
              file, str(dim), str(iterations)]


def benchmark_all_mtx(files, config, mtx_folder, dense_range, iterations = 100, findbest=True):
   if "conf_keys" in dir(config):
      print((f"|mtx|{config.baseline_key}|"+"{}|"*len(config.conf_keys)).format(*config.conf_keys))
      print("|--|:--:|"+":--:|"*len(config.conf_keys))

   for f in files:
      for dim in dense_range:
         failed = False
         # exec_cmd = config._launch_args(gen_file_path(mtx_folder, f), dim, iterations)
         exec_cmd = config._launch_args(gen_sample_path(mtx_folder, f), dim, iterations)
         try:
            out = subp.check_output(exec_cmd)
         except subp.CalledProcessError:
            failed = True
            print(f"running {exec_cmd} failed")
         
         if not failed:
            out = str(out)[2:-1]
            out = out.strip().split("\\n")

            cu_times = [config._find_baseline_time(x) for x in out]
            my_times = [config._find_my_time(x) for x in out]

            cu_time = max(cu_times)
            my_time = [x for x in my_times if x > 0]
            if findbest:
               my_best = min(my_time)
               my_conf = config._get_config_string(out[my_times.index(my_best)])
               print(f"{f}-{dim} my={my_time} cu={cu_time} speedup {cu_time/my_time} Best conf {my_conf}")
            else:
               if "conf_keys" in dir(config):
                  keys = config.conf_keys
                  row = ["/"]*(len(keys))
                  for i in range(len(my_times)):
                     i_time = my_times[i]
                     if i_time < 0:
                        continue
                     i_conf = config._get_config_string(out[i])
                     row[keys.index(i_conf)] = i_time
                  row = [cu_time] + row
                  print((f"|{f}-{dim}" + "|{:.6}"*len(row) + "|").format(*row), flush=True)

               else:
                  print(f"{f}-{dim} baseline = {cu_time}")
                  for i in range(len(my_times)):
                     i_time = my_times[i]
                     if i_time < 0:
                        continue
                     i_conf = config._get_config_string(out[i])
                     print(f"{i_conf} = {i_time}")
                  print("----------------------")
         # bench_time = {}
         # bench_speedup = ""
         # for line in times[1:]:
         #    line = line.strip().split(" ")
         #    bench_time[line[0]] = float(line[1])
         
         # for k,v in bench_time.items():
         #    if v == min(bench_time.values()):
         #       bench_speedup += "{} : \033[1;31m{:.6f}\033[0m, ".format(k,base_time/v)
         #    else:
         #       bench_speedup += "{} : {:.6f}, ".format(k,base_time/v)

         # print("{} dim[{}]\n time {}\n speedup {}".format(f, dim, bench_time, bench_speedup)) 

Sampling = os.listdir("./data/Mtx/raw")

if __name__ == '__main__':
   import sys
   print(sys.argv)
   assert(len(sys.argv) >= 3)
   param = sys.argv[1:]
   se = [int(x) for x in param[1].split('~')]
   i = se[-1]//2
   while i > se[0]:
      se.insert(1, i)
      i //= 2
   param[1] = se
   # run_all_mtx(*param)
   # benchmark_all_mtx(Sampling, GeSpMMMotiveConfig, *param, iterations=200, findbest=False)
   # benchmark_all_mtx(Sampling, MyConfig, *param, iterations=50, findbest=False)
   benchmark_all_mtx(Sampling, ManscriptConfig, *param, iterations=50, findbest=False)